# $Id: session.py,v 1.14 2010/03/20 18:29:22 jribbens Exp $

import time, hmac, re, random, os, errno
try:
    import hashlib
    sha = hashlib.sha1
    shanew = hashlib.sha1
except ImportError:
    import sha
    shanew = sha.new
try:
    import cPickle as pickle
except ImportError:
    import pickle

class Error(Exception):
  pass

class Session(dict):
  def _make_hash(self, sid, secret):
    """Create a hash for 'sid'
    
    This function may be overridden by subclasses."""
    return hmac.new(secret, sid, sha).hexdigest()[:8]

  def _create(self, secret):
    """Create a new session ID and, optionally hash
    
    This function must insert the new session ID (which must be 8 hexadecimal
    characters) into self["id"].

    It may optionally insert the hash into self["hash"]. If it doesn't, then
    _make_hash will automatically be called later.

    This function may be overridden by subclasses.
    """
    rnd = str(time.time()) + str(random.random()) + \
      str(self._req["UNIQUE_ID"])
    self["id"] = shanew(rnd).hexdigest()[:8]

  def _load(self):
    """Load the session dictionary from somewhere
    
    This function may be overridden by subclasses.
    It should return 1 if the load was successful, or 0 if the session could
    not be found. Any other type of error should raise an exception as usual."""
    return 1

  def save(self):
    """Save the session dictionary to somewhere
    
    This function may be overridden by subclasses."""
    pass

  def tidy():
    pass
  tidy = staticmethod(tidy)

  def __init__(self, req, secret, root="", sid=None, shash=None, secure=0, domain=None, create_new=True):
    dict.__init__(self)
    self["id"] = None
    self._req = req
    self.relocated = 0
    self.new = 0
    
    # try and determine existing session id
    if sid is not None:
      self["id"] = sid
      if shash is None:
        self["hash"] = self._make_hash(self["id"], secret)
      else:
        self["hash"] = shash
        if self["hash"] != self._make_hash(self["id"], secret):
          self["id"] = None

    # try and load the session
    if self["id"] is not None:
      if not self._load():
        self["id"] = None

    # if no session was available and loaded, create a new one
    if self["id"] is None and create_new:
      if "hash" in self:
        del self["hash"]
      self.created = time.time()
      self.new = 1
      self._create(secret)
      if "hash" not in self:
        self["hash"] = self._make_hash(self["id"], secret)


class MemorySession(Session):
  _sessions = {}

  def _create(self, secret):
    while 1:
      Session._create(self, secret)
      if self["id"] in self._sessions:
        continue
      self._sessions[self["id"]] = {"created": self.created,
        "updated": self.created, "data": {}}
      break
  
  def _load(self):
    try:
      sess = self._sessions[self["id"]]
    except KeyError:
      return 0
    self.created = sess["created"]
    self.update(sess["data"])
    return 1

  def save(self):
    sess = self._sessions[self["id"]]
    sess["updated"] = time.time()
    sess["data"] = self.copy()

  def tidy(cls, max_idle=0, max_age=0):
    now = time.time()
    for k in cls._sessions:
      if (max_age and k["created"] < now - max_age) or \
        (max_idle and k["updated"] < now - max_idle):
        del cls._sessions[k]
  tidy = classmethod(tidy)

class GenericSQLSession(Session):
    def _create(self, secret):
        while 1:
            Session._create(self, secret)
            self["hash"] = self._make_hash(self["id"], secret)
            try:
                self.dbc.execute("INSERT INTO %s (ID,hash,created,updated,data)"
                    " VALUES (%%s,%%s,%%s,%%s,%%s)" % (self.table,),
                    (self["id"], self["hash"], int(self.created), int(self.created),
                    pickle.dumps({}, 1)))
                self.dbc.execute("COMMIT")
            except self.dbc.IntegrityError:
                pass
            else:
                break

    def _load(self):
        self.dbc.execute("SELECT created,data FROM %s WHERE ID=%%s" % (self.table,),
            (self["id"],))
        if self.dbc.rowcount == 0:
            return 0
        row = self.dbc.fetchone()
        self.created = row[0]
        self.update(pickle.loads(row[1]))
        return 1

    def save(self):
        self.dbc.execute("UPDATE %s SET updated=%%s,data=%%s"
            " WHERE ID=%%s" % (self.table,), (int(time.time()),
            pickle.dumps(self.copy(), 1), self["id"]))
        self.dbc.execute("COMMIT")

    @staticmethod
    def tidy(dbc, table="sessions", max_idle=0, max_age=0):
        now = time.time()
        if max_idle:
            dbc.execute("DELETE FROM %s WHERE updated < %%s" % (table,),
                (now - max_idle,))
        if max_age:
            dbc.execute("DELETE FROM %s WHERE created < %%s" % (table,),
                (now - max_age,))
        if max_idle or max_age:
            dbc.execute("COMMIT")
    
    def remove(self):
        self.dbc.execute("DELETE FROM %s WHERE ID = %%s AND hash = %%s" % (self.table,), (long(self['id'], 16), self['hash'],))
        self.dbc.execute("COMMIT")

    def __init__(self, req, secret, dbc, table="sessions", **kwargs):
        self.dbc = dbc
        self.table = table
        Session.__init__(self, req, secret, **kwargs)

class MySQLSession(GenericSQLSession):
    def _create(self, secret):
        self.dbc.execute("LOCK TABLES %s WRITE" % (self.table,))
        while 1:
            Session._create(self, secret)
            self.dbc.execute("SELECT 1 FROM %s WHERE ID=%%s" % (self.table,),
                (long(self["id"], 16),))
            if self.dbc.rowcount == 0:
                break
        self["hash"] = self._make_hash(self["id"], secret)
        self.dbc.execute("INSERT INTO %s (ID,hash,created,updated,data) VALUES " \
            "(%%s,%%s,%%s,%%s,%%s)" % (self.table,),
            (long(self["id"], 16), self["hash"], int(self.created),
            int(self.created), pickle.dumps({}, 1)))
        self.dbc.execute("UNLOCK TABLES")

    def _load(self):
        self.dbc.execute("SELECT created,data FROM %s WHERE ID=%%s" % (self.table,),
            (long(self["id"], 16),))
        if self.dbc.rowcount == 0:
            return 0
        row = self.dbc.fetchone()
        self.created = row[0]
        self.update(pickle.loads(row[1]))
        return 1

    def save(self):
        serialized = pickle.dumps(self.copy(), 1)
        self.dbc.execute("UPDATE %s SET updated=%%s,data=%%s WHERE ID=%%s" % (self.table,), (int(time.time()), serialized, long(self["id"], 16)))
        self.dbc.execute("COMMIT")

SQLSession = MySQLSession # backwards compatibility name
